package cn.eleven.util;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.sql.*;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import org.java_websocket.server.WebSocketServer;
import org.java_websocket.WebSocket;
import org.java_websocket.handshake.ClientHandshake;

public class ChatServer extends WebSocketServer {
    private static final int PORT = 8080;
    private static final Map<String, WebSocket> clients = Collections.synchronizedMap(new HashMap<>());

    public ChatServer() {
        super(new InetSocketAddress(PORT));
    }

    @Override
    public void onOpen(WebSocket conn, ClientHandshake handshake) {
        String username = handshake.getResourceDescriptor().replace("/?username=", "");
        clients.put(username, conn);
        System.out.println(username + " 连接到服务器");
    }

    @Override
    public void onMessage(WebSocket conn, String message) {
        String[] parts = message.split("\\|", 2);
        if (parts.length == 2) {
            String receiver = parts[0];
            String msg = parts[1];
            String sender = getUsername(conn);

            saveMessage(sender, receiver, msg);

            WebSocket recipientConn = clients.get(receiver);
            if (recipientConn != null && recipientConn.isOpen()) {
                recipientConn.send(sender + ": " + msg);
            }
        }
    }

    @Override
    public void onClose(WebSocket conn, int code, String reason, boolean remote) {
        String username = getUsername(conn);
        clients.remove(username);
        System.out.println(username + " 断开连接");
    }

    @Override
    public void onError(WebSocket conn, Exception ex) {
        ex.printStackTrace();
    }

    @Override
    public void onStart() {

    }

    private String getUsername(WebSocket conn) {
        for (Map.Entry<String, WebSocket> entry : clients.entrySet()) {
            if (entry.getValue().equals(conn)) {
                return entry.getKey();
            }
        }
        return null;
    }

    private void saveMessage(String sender, String receiver, String message) {
        String url = "jdbc:mysql://localhost:3306/chat_db";
        String user = "root";
        String password = "1234";

        try (Connection conn = DriverManager.getConnection(url, user, password);
             PreparedStatement stmt = conn.prepareStatement(
                     "INSERT INTO messages (sender, receiver, message) VALUES (?, ?, ?)")) {
            stmt.setString(1, sender);
            stmt.setString(2, receiver);
            stmt.setString(3, message);
            stmt.executeUpdate();
        } catch (SQLException e) {
            e.printStackTrace();
        }
    }

    public static void main(String[] args) {
        new ChatServer().start();
        System.out.println("WebSocket 服务器启动，端口：" + PORT);
    }
}
